import * as tf from "@tensorflow/tfjs-node";

function createModel(): tf.Sequential {
  const model: tf.Sequential = tf.sequential();
  model.add(
    tf.layers.dense({
      units: 1,
      useBias: true,
      activation: "linear",
      inputDim: 1,
    })
  );

  const optim: tf.SGDOptimizer = tf.train.sgd(0.1);
  // const optim: tf.AdamOptimizer = tf.train.adam(0.1);
  // 在这里，Adam没有SGD好

  model.compile({
    loss: "meanSquaredError",
    optimizer: optim,
  });

  return model;
}

export { createModel };
